"""
数据库结构提取器
负责提取数据库的完整结构信息，包括表、字段、约束等
"""

import logging
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, asdict
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

from .connection_manager import db_manager
from config.database_config import DatabaseType


@dataclass
class ColumnInfo:
    """列信息"""
    column_name: str
    data_type: str
    nullable: bool
    default_value: Optional[str] = None
    is_primary_key: bool = False
    is_foreign_key: bool = False
    auto_increment: bool = False
    column_comment: str = ""
    character_maximum_length: Optional[int] = None
    numeric_precision: Optional[int] = None
    numeric_scale: Optional[int] = None
    
    def __post_init__(self):
        # 确保注释字段不为None
        if self.column_comment is None:
            self.column_comment = ""


@dataclass
class ForeignKeyInfo:
    """外键信息"""
    name: str
    column_name: str
    referenced_table: str
    referenced_column: str
    referenced_schema: Optional[str] = None


@dataclass
class IndexInfo:
    """索引信息"""
    name: str
    columns: List[str]
    is_unique: bool = False
    is_primary: bool = False


@dataclass
class TableInfo:
    """表信息"""
    table_name: str
    table_comment: str = ""
    schema_name: Optional[str] = None
    columns: List[ColumnInfo] = None
    foreign_keys: List[ForeignKeyInfo] = None
    indexes: List[IndexInfo] = None
    row_count: Optional[int] = None
    table_size: Optional[str] = None
    
    def __post_init__(self):
        if self.columns is None:
            self.columns = []
        if self.foreign_keys is None:
            self.foreign_keys = []
        if self.indexes is None:
            self.indexes = []
        # 确保注释字段不为None
        if self.table_comment is None:
            self.table_comment = ""
        # 确保所有列的注释字段不为None
        for column in self.columns:
            if column.column_comment is None:
                column.column_comment = ""


@dataclass
class DatabaseSchema:
    """数据库结构"""
    database_name: str
    config_name: str
    database_type: DatabaseType
    schemas: Dict[str, List[TableInfo]] = None
    extraction_time: Optional[str] = None
    total_tables: int = 0
    total_columns: int = 0
    
    def __post_init__(self):
        if self.schemas is None:
            self.schemas = {}


class SchemaExtractor:
    """数据库结构提取器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.lock = threading.RLock()
    
    def extract_database_schema(self, config_name: str, include_data_info: bool = True, 
                              schema_filter: Optional[List[str]] = None,
                              table_filter: Optional[List[str]] = None,
                              max_workers: int = 5) -> Optional[DatabaseSchema]:
        """
        提取完整的数据库结构
        
        Args:
            config_name: 数据库配置名称
            include_data_info: 是否包含数据信息（行数、大小等）
            schema_filter: 指定要提取的schema列表
            table_filter: 指定要提取的表列表
            max_workers: 最大并发工作线程数
        """
        try:
            # 测试连接
            if not db_manager.test_connection(config_name):
                self.logger.error(f"数据库连接失败：{config_name}")
                return None
            
            # 获取数据库配置
            from config.database_config import db_config_manager
            db_config = db_config_manager.get_config(config_name)
            if not db_config:
                self.logger.error(f"数据库配置不存在：{config_name}")
                return None
            
            # 创建数据库结构对象
            schema_obj = DatabaseSchema(
                database_name=db_config.database,
                config_name=config_name,
                database_type=db_config.db_type
            )
            
            # 获取所有schema
            schema_names = self._get_schemas(config_name, schema_filter)
            
            self.logger.info(f"开始提取数据库结构：{config_name}，共 {len(schema_names)} 个schema")
            
            # 并发提取每个schema的表结构
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                future_to_schema = {
                    executor.submit(self._extract_schema_tables, config_name, schema_name, 
                                  table_filter, include_data_info): schema_name 
                    for schema_name in schema_names
                }
                
                for future in as_completed(future_to_schema):
                    schema_name = future_to_schema[future]
                    try:
                        tables = future.result()
                        if tables:
                            schema_obj.schemas[schema_name] = tables
                            schema_obj.total_tables += len(tables)
                            schema_obj.total_columns += sum(len(table.columns) for table in tables)
                        
                        self.logger.info(f"Schema {schema_name} 提取完成，共 {len(tables) if tables else 0} 个表")
                        
                    except Exception as e:
                        self.logger.error(f"提取schema {schema_name} 失败: {str(e)}")
            
            # 设置提取时间
            from datetime import datetime
            schema_obj.extraction_time = datetime.now().isoformat()
            
            self.logger.info(f"数据库结构提取完成：{config_name}，"
                           f"共 {schema_obj.total_tables} 个表，{schema_obj.total_columns} 个字段")
            
            return schema_obj
            
        except Exception as e:
            self.logger.error(f"提取数据库结构失败 {config_name}: {str(e)}")
            return None
    
    def _get_schemas(self, config_name: str, schema_filter: Optional[List[str]]) -> List[str]:
        """获取schema列表"""
        try:
            from config.database_config import db_config_manager
            db_config = db_config_manager.get_config(config_name)
            
            # SQLite没有schema概念，直接返回空列表表示主数据库
            if db_config.db_type == DatabaseType.SQLITE:
                return ["main"]
            
            schemas = db_manager.get_schema_names(config_name)
            
            # 过滤系统schema
            system_schemas = {
                DatabaseType.MYSQL: ["information_schema", "performance_schema", "mysql", "sys"],
                DatabaseType.POSTGRESQL: ["information_schema", "pg_catalog", "pg_toast"],
                DatabaseType.SQLSERVER: ["information_schema", "sys", "INFORMATION_SCHEMA"],
                DatabaseType.ORACLE: ["SYS", "SYSTEM", "SYSAUX", "CTXSYS", "MDSYS", "OLAPSYS", "ORDSYS", "OUTLN", "WMSYS"]
            }
            
            if db_config.db_type in system_schemas:
                schemas = [s for s in schemas if s not in system_schemas[db_config.db_type]]
            
            # 应用schema过滤器
            if schema_filter:
                schemas = [s for s in schemas if s in schema_filter]
            
            # 如果没有指定schema或过滤后为空，则使用默认schema
            if not schemas:
                if db_config.db_type == DatabaseType.MYSQL:
                    schemas = [db_config.database]
                elif db_config.db_type == DatabaseType.POSTGRESQL:
                    schemas = ["public"]
                elif db_config.db_type == DatabaseType.SQLSERVER:
                    schemas = ["dbo"]
                else:
                    schemas = [None]  # Oracle等使用None表示当前用户的schema
            
            return schemas
            
        except Exception as e:
            self.logger.error(f"获取schema列表失败: {str(e)}")
            return []
    
    def _extract_schema_tables(self, config_name: str, schema_name: str, 
                             table_filter: Optional[List[str]], 
                             include_data_info: bool) -> List[TableInfo]:
        """提取指定schema的所有表结构"""
        try:
            # 获取表列表
            table_names = db_manager.get_table_names(config_name, schema=schema_name if schema_name != "main" else None)
            
            # 应用表过滤器
            if table_filter:
                table_names = [t for t in table_names if t in table_filter]
            
            if not table_names:
                return []
            
            tables = []
            
            # 提取每个表的详细信息
            for table_name in table_names:
                try:
                    table_info = self._extract_table_info(config_name, table_name, schema_name, include_data_info)
                    if table_info:
                        tables.append(table_info)
                except Exception as e:
                    self.logger.error(f"提取表 {schema_name}.{table_name} 失败: {str(e)}")
            
            return tables
            
        except Exception as e:
            self.logger.error(f"提取schema {schema_name} 的表失败: {str(e)}")
            return []
    
    def _extract_table_info(self, config_name: str, table_name: str, schema_name: str, 
                          include_data_info: bool) -> Optional[TableInfo]:
        """提取单个表的详细信息"""
        try:
            # 获取基本表信息
            table_detail = db_manager.get_table_info(config_name, table_name, 
                                                   schema=schema_name if schema_name != "main" else None)
            if not table_detail:
                return None
            
            # 创建表信息对象
            table_info = TableInfo(
                table_name=table_name,
                schema_name=schema_name,
                table_comment=table_detail.get("table_comment") or ""
            )
            
            # 处理列信息
            pk_constraint = table_detail.get("primary_keys", {})
            primary_keys = set(pk_constraint.get("constrained_columns", []) if pk_constraint else [])
            foreign_key_columns = {fk["constrained_columns"][0]: fk for fk in table_detail.get("foreign_keys", []) 
                                 if fk.get("constrained_columns")}
            
            for col in table_detail.get("columns", []):
                column_info = ColumnInfo(
                    column_name=col["name"],
                    data_type=str(col["type"]),
                    nullable=col["nullable"],
                    default_value=str(col["default"]) if col["default"] is not None else None,
                    is_primary_key=col["name"] in primary_keys,
                    is_foreign_key=col["name"] in foreign_key_columns,
                    auto_increment=col.get("autoincrement", False),
                    column_comment=col.get("comment") or ""
                )
                table_info.columns.append(column_info)
            
            # 处理外键信息
            for fk in table_detail.get("foreign_keys", []):
                if fk.get("constrained_columns") and fk.get("referred_columns"):
                    fk_info = ForeignKeyInfo(
                        name=fk.get("name", ""),
                        column_name=fk["constrained_columns"][0],
                        referenced_table=fk.get("referred_table", ""),
                        referenced_column=fk["referred_columns"][0],
                        referenced_schema=fk.get("referred_schema")
                    )
                    table_info.foreign_keys.append(fk_info)
            
            # 处理索引信息
            for idx in table_detail.get("indexes", []):
                idx_info = IndexInfo(
                    name=idx.get("name", ""),
                    columns=idx.get("column_names", []),
                    is_unique=idx.get("unique", False),
                    is_primary=idx.get("name", "").lower().startswith("primary")
                )
                table_info.indexes.append(idx_info)
            
            # 获取数据信息（如果需要）
            if include_data_info:
                try:
                    table_info.row_count = self._get_table_row_count(config_name, table_name, schema_name)
                    table_info.table_size = self._get_table_size(config_name, table_name, schema_name)
                except Exception as e:
                    self.logger.warning(f"获取表 {table_name} 数据信息失败: {str(e)}")
            
            return table_info
            
        except Exception as e:
            self.logger.error(f"提取表信息失败 {table_name}: {str(e)}")
            return None
    
    def _get_table_row_count(self, config_name: str, table_name: str, schema_name: str) -> Optional[int]:
        """获取表行数"""
        try:
            schema_prefix = f"{schema_name}." if schema_name and schema_name != "main" else ""
            query = f"SELECT COUNT(*) as row_count FROM {schema_prefix}{table_name}"
            result = db_manager.fetch_one(config_name, query)
            return result["row_count"] if result else None
        except Exception as e:
            self.logger.warning(f"获取表行数失败 {table_name}: {str(e)}")
            return None
    
    def _get_table_size(self, config_name: str, table_name: str, schema_name: str) -> Optional[str]:
        """获取表大小（这个方法需要根据不同数据库类型实现）"""
        try:
            from config.database_config import db_config_manager
            db_config = db_config_manager.get_config(config_name)
            if not db_config:
                return None
            
            # 不同数据库获取表大小的方法不同
            if db_config.db_type == DatabaseType.MYSQL:
                query = """
                SELECT ROUND(((data_length + index_length) / 1024 / 1024), 2) AS size_mb
                FROM information_schema.tables 
                WHERE table_schema = %s AND table_name = %s
                """
                result = db_manager.fetch_one(config_name, query, {
                    "table_schema": db_config.database,
                    "table_name": table_name
                })
                return f"{result['size_mb']} MB" if result and result.get('size_mb') else None
            
            elif db_config.db_type == DatabaseType.POSTGRESQL:
                query = "SELECT pg_size_pretty(pg_total_relation_size(%s)) AS size"
                full_table_name = f"{schema_name}.{table_name}" if schema_name and schema_name != "public" else table_name
                result = db_manager.fetch_one(config_name, query, {"table_name": full_table_name})
                return result["size"] if result else None
            
            # 其他数据库类型可以后续扩展
            return None
            
        except Exception as e:
            self.logger.warning(f"获取表大小失败 {table_name}: {str(e)}")
            return None


# 全局结构提取器实例
schema_extractor = SchemaExtractor()